import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import os

# 设置中文字体支持
def set_chinese_font():
    """设置中文字体"""
    try:
        if os.name == 'nt':  # Windows系统
            matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
        else:  # Linux/Mac系统
            matplotlib.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei', 'DejaVu Sans']
        
        matplotlib.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
        return True
    except:
        print("警告: 中文字体设置失败，将使用默认字体")
        return False

# 设置随机种子以保证结果可重现
tf.random.set_seed(42)
np.random.seed(42)

class MNISTClassifier:
    def __init__(self, binarize_threshold=0.5):
        self.model = None
        self.history = None
        self.binarize_threshold = binarize_threshold
        self.x_train = None  # 保存训练数据用于量化
    
    def binarize_image(self, image):
        """将图像二值化"""
        binary_image = np.where(image > self.binarize_threshold, 1.0, 0.0)
        return binary_image
        
    def load_data(self):
        """加载MNIST数据集"""
        print("正在加载MNIST数据集...")
        (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
        
        # 数据预处理
        x_train = x_train.astype('float32') / 255.0
        x_test = x_test.astype('float32') / 255.0
        
        # 重塑为(28,28,1)
        x_train = x_train.reshape(x_train.shape[0], 28, 28, 1)
        x_test = x_test.reshape(x_test.shape[0], 28, 28, 1)
        
        # 标签one-hot编码
        y_train = keras.utils.to_categorical(y_train, 10)
        y_test = keras.utils.to_categorical(y_test, 10)
        
        # 二值化处理
        print("正在进行二值化处理...")
        x_train = self.binarize_image(x_train)
        x_test = self.binarize_image(x_test)
        
        self.x_train = x_train  # 保存训练数据用于量化
        return (x_train, y_train), (x_test, y_test)
    
    def build_model(self):
        """构建CNN模型"""
        print("正在构建模型...")
        self.model = keras.Sequential([
            # 第一个卷积层
            keras.layers.Conv2D(16, (3, 3), activation='relu', input_shape=(28, 28, 1)),
            keras.layers.MaxPooling2D((2, 2)),
            keras.layers.BatchNormalization(),
         
            keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same'),
            keras.layers.BatchNormalization(),
            keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same'),
            keras.layers.BatchNormalization(),
            keras.layers.MaxPooling2D(),
            keras.layers.Dropout(0.25),

            keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
            keras.layers.BatchNormalization(),
            keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
            keras.layers.BatchNormalization(),
            keras.layers.MaxPooling2D(),
            keras.layers.Dropout(0.25),
            
            # 展平层
            keras.layers.Flatten(),
            
            # 全连接层
            keras.layers.Dense(64, activation='relu'),
            keras.layers.Dropout(0.5),
            
            # 输出层
            keras.layers.Dense(10, activation='softmax')
        ])
        
        # 编译模型
        self.model.compile(optimizer='adam',
                          loss='categorical_crossentropy',
                          metrics=['accuracy'])
        
        print("模型结构:")
        self.model.summary()
    
    def train_model(self, x_train, y_train, x_test, y_test, epochs=5, batch_size=128):
        """训练模型"""
        print("开始训练模型...")
        
        # 设置回调函数
        callbacks = [
            keras.callbacks.EarlyStopping(patience=3, restore_best_weights=True),
            keras.callbacks.ReduceLROnPlateau(factor=0.2, patience=2)
        ]
        
        # 训练模型
        self.history = self.model.fit(
            x_train, y_train,
            batch_size=batch_size,
            epochs=epochs,
            validation_data=(x_test, y_test),
            callbacks=callbacks,
            verbose=1
        )
    
    def evaluate_model(self, x_test, y_test):
        """评估模型性能"""
        print("评估模型性能...")
        test_loss, test_acc = self.model.evaluate(x_test, y_test, verbose=0)
        print(f"测试集准确率: {test_acc:.4f}")
        print(f"测试集损失: {test_loss:.4f}")
        return test_loss, test_acc
    
    def plot_training_history(self):
        """绘制训练历史"""
        if self.history is None:
            print("没有训练历史可显示")
            return
            
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
        
        # 绘制准确率曲线
        ax1.plot(self.history.history['accuracy'], label='训练准确率')
        ax1.plot(self.history.history['val_accuracy'], label='验证准确率')
        ax1.set_title('模型准确率')
        ax1.set_xlabel('训练轮次')
        ax1.set_ylabel('准确率')
        ax1.legend()
        ax1.grid(True)
        
        # 绘制损失曲线
        ax2.plot(self.history.history['loss'], label='训练损失')
        ax2.plot(self.history.history['val_loss'], label='验证损失')
        ax2.set_title('模型损失')
        ax2.set_xlabel('训练轮次')
        ax2.set_ylabel('损失值')
        ax2.legend()
        ax2.grid(True)
        
        plt.tight_layout()
        plt.show()
    
    def predict_single_image(self, image):
        """预测单个图像"""
        # 确保图像形状正确
        if len(image.shape) == 2:
            image = image.reshape(1, 28, 28, 1)
        elif len(image.shape) == 3 and image.shape[-1] != 1:
            if image.shape[-1] == 3:
                image = np.dot(image[...,:3], [0.2989, 0.5870, 0.1140])
                image = image.reshape(1, 28, 28, 1)
            else:
                image = image.reshape(1, 28, 28, 1)
        elif len(image.shape) == 3 and image.shape[-1] == 1:
            image = image.reshape(1, 28, 28, 1)
        
        # 确保数据类型正确
        image = image.astype('float32')
        
        # 二值化处理
        image = self.binarize_image(image)
        
        prediction = self.model.predict(image, verbose=0)
        predicted_class = np.argmax(prediction)
        confidence = np.max(prediction)
        
        return predicted_class, confidence, prediction
    
    def display_prediction(self, x_test, y_test, num_samples=5):
        """显示预测结果示例"""
        print("显示预测示例...")
        
        indices = np.random.choice(len(x_test), num_samples)
        
        fig, axes = plt.subplots(1, num_samples, figsize=(15, 3))
        if num_samples == 1:
            axes = [axes]
        
        for i, idx in enumerate(indices):
            image = x_test[idx]
            true_label = np.argmax(y_test[idx])
            
            if len(image.shape) == 3 and image.shape[2] == 1:
                display_image = image.reshape(28, 28)
            else:
                display_image = image
            
            pred_class, confidence, _ = self.predict_single_image(image)
            
            axes[i].imshow(display_image, cmap='gray')
            axes[i].set_title(f'真实: {true_label}\n预测: {pred_class}\n置信度: {confidence:.2f}')
            axes[i].axis('off')
            
            color = 'green' if true_label == pred_class else 'red'
            for spine in axes[i].spines.values():
                spine.set_color(color)
                spine.set_linewidth(3)
        
        plt.tight_layout()
        plt.show()

    def convert_to_tflite_int8(self, output_path='mnist_cnn_model_int8.tflite'):
        """将模型转换为输入输出都是int8类型的TFLite格式"""
        if self.model is None:
            print("模型尚未初始化，无法转换为TFLite格式")
            return
            
        print(f"正在将模型转换为int8输入输出的TFLite格式...")
        
        # 创建TFLite转换器
        converter = tf.lite.TFLiteConverter.from_keras_model(self.model)
        
        # 定义代表性数据集生成器（用于量化校准）
        def representative_dataset():
            for i in range(100):
                # 取训练数据样本并转换为float32
                sample = tf.cast(self.x_train[i:i+1], tf.float32)
                yield [sample]
        
        # 配置全整数量化参数
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
        converter.representative_dataset = representative_dataset
        # 设置支持的操作集为INT8
        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
        # 明确指定输入输出类型为INT8
        converter.inference_input_type = tf.int8
        converter.inference_output_type = tf.int8
        
        # 转换模型
        tflite_model = converter.convert()
        
        # 保存TFLite模型
        with open(output_path, 'wb') as f:
            f.write(tflite_model)
        
        print(f"INT8输入输出的TFLite模型已保存为: {output_path}")

def main():
    """主函数"""
    # 设置中文字体
    set_chinese_font()
    
    print("TensorFlow手写数字识别示例（二值图像）")
    print("=" * 50)
    
    # 创建分类器实例
    classifier = MNISTClassifier(binarize_threshold=0.5)
    
    # 加载数据
    (x_train, y_train), (x_test, y_test) = classifier.load_data()
    
    print(f"训练集形状: {x_train.shape}")
    print(f"测试集形状: {x_test.shape}")
    print(f"训练集像素值范围: [{x_train.min()}, {x_train.max()}]")
    print(f"测试集像素值范围: [{x_test.min()}, {x_test.max()}]")
    
    # 构建模型
    classifier.build_model()
    
    # 训练模型
    classifier.train_model(x_train, y_train, x_test, y_test, epochs=5)
    
    # 评估模型
    classifier.evaluate_model(x_test, y_test)
    
    # 显示训练历史
    classifier.plot_training_history()
    
    # 显示预测示例
    classifier.display_prediction(x_test, y_test, num_samples=5)
    
    print("程序执行完成！")
    classifier.model.save('mnist_cnn_model_binary.h5')
    print("模型已保存为 'mnist_cnn_model_binary.h5'")
    
    # 转换为INT8 TFLite模型
    classifier.convert_to_tflite_int8()

if __name__ == "__main__":
    main()